import pymysql
from pandas import DataFrame
from pymysql.connections import Connection
from sqlalchemy import create_engine, Engine


# 定义一个函数，用来创建连接数据库
def mysql_db() -> Connection:
    """

    :rtype: Connection
    """
    return pymysql.connect(host='localhost',
                           port=3306,
                           database='world',
                           user='root',
                           password='zxc!110258',
                           charset='utf8')


def engine() -> Engine:
    """
    :rtype: engine
    """
    return create_engine('mysql+pymysql://root:zxc!110258@localhost/world')


def query(sql: str) -> DataFrame:
    with mysql_db() as con:
        with con.cursor() as cursor:
            cursor.execute(sql)
            columns = dict()
            for index in range(len(cursor.description)):
                columns.update({index: cursor.description[index][0]})
            data = DataFrame(data=cursor.fetchall())
            data.rename(columns=columns, inplace=True)
            return data


def update(sql: str) -> int:
    """

    :rtype: object
    """
    with mysql_db() as con:
        with con.cursor() as cursor:
            result = cursor.execute(sql)
            con.commit()
            return result
